﻿"""
人脸识别模型
处理人脸检测、特征提取和识别
"""

import cv2
import numpy as np
from typing import Dict, Tuple, Optional, List
from src.config.settings import FACE_RECOGNITION_CONFIG


class FaceRecognitionModel:
    """人脸识别模型"""
    
    def __init__(self):
        self.known_faces = {}  # 存储已知人脸信息
        self.face_cascade = cv2.CascadeClassifier(
            cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
        )
        self.recognition_threshold = FACE_RECOGNITION_CONFIG['recognition_threshold']
        self.min_face_size = FACE_RECOGNITION_CONFIG['min_face_size']
    
    def extract_face_features(self, face_image: np.ndarray) -> Optional[np.ndarray]:
        """提取人脸特征"""
        try:
            # 转换为灰度图
            gray = cv2.cvtColor(face_image, cv2.COLOR_BGR2GRAY)
            
            # 调整大小
            resized = cv2.resize(gray, self.min_face_size)
            
            # 直方图均衡化，增强对比度
            equalized = cv2.equalizeHist(resized)
            
            # 提取LBP特征（局部二值模式）
            lbp = self.extract_lbp_features(equalized)
            
            # 提取HOG特征（方向梯度直方图）
            hog = self.extract_hog_features(equalized)
            
            # 组合特征
            features = np.concatenate([lbp, hog])
            
            # 归一化
            features = features / np.linalg.norm(features)
            
            return features
            
        except Exception as e:
            print(f"提取人脸特征失败: {e}")
            return None
    
    def extract_lbp_features(self, gray_image: np.ndarray) -> np.ndarray:
        """提取LBP特征"""
        try:
            # 检查输入图像
            if gray_image is None or gray_image.size == 0:
                return np.zeros(256)
            
            height, width = gray_image.shape
            if height < 3 or width < 3:
                return np.zeros(256)
            
            # 简单的LBP实现
            lbp_image = np.zeros_like(gray_image)
            
            # 使用向量化操作提高性能，避免嵌套循环
            for i in range(1, height - 1):
                for j in range(1, width - 1):
                    center = gray_image[i, j]
                    code = 0
                    code |= (gray_image[i - 1, j - 1] >= center) << 7
                    code |= (gray_image[i - 1, j] >= center) << 6
                    code |= (gray_image[i - 1, j + 1] >= center) << 5
                    code |= (gray_image[i, j + 1] >= center) << 4
                    code |= (gray_image[i + 1, j + 1] >= center) << 3
                    code |= (gray_image[i + 1, j] >= center) << 2
                    code |= (gray_image[i + 1, j - 1] >= center) << 1
                    code |= (gray_image[i, j - 1] >= center) << 0
                    lbp_image[i, j] = code
            
            # 计算直方图作为特征
            hist, _ = np.histogram(lbp_image.ravel(), bins=256, range=[0, 256])
            hist_sum = np.sum(hist)
            if hist_sum > 0:
                return hist / hist_sum  # 归一化
            else:
                return np.zeros(256)
                
        except Exception as e:
            print(f"提取LBP特征失败: {e}")
            return np.zeros(256)
    
    def extract_hog_features(self, gray_image: np.ndarray) -> np.ndarray:
        """提取HOG特征"""
        try:
            # 检查输入图像
            if gray_image is None or gray_image.size == 0:
                return np.zeros(9)
            
            # 计算梯度
            gx = cv2.Sobel(gray_image, cv2.CV_32F, 1, 0)
            gy = cv2.Sobel(gray_image, cv2.CV_32F, 0, 1)
            
            # 计算梯度幅值和方向
            mag, angle = cv2.cartToPolar(gx, gy, angleInDegrees=True)
            
            # 简化的HOG特征（9个方向的直方图）
            hist, _ = np.histogram(angle, bins=9, range=[0, 180], weights=mag)
            hist_sum = np.sum(hist)
            if hist_sum > 0:
                return hist / hist_sum  # 归一化
            else:
                return np.zeros(9)
                
        except Exception as e:
            print(f"提取HOG特征失败: {e}")
            return np.zeros(9)
    
    def compare_faces(self, features1: np.ndarray, features2: np.ndarray) -> float:
        """比较两个人脸特征的相似度"""
        if features1 is None or features2 is None:
            return 0.0
        
        # 计算余弦相似度
        similarity = np.dot(features1, features2) / (
            np.linalg.norm(features1) * np.linalg.norm(features2) + 1e-10
        )
        return max(0.0, min(1.0, similarity))
    
    def detect_faces(self, frame: np.ndarray) -> List[Tuple[int, int, int, int]]:
        """检测人脸"""
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(
            gray, 1.1, 4, minSize=self.min_face_size
        )
        return faces.tolist()
    
    def recognize_face(self, face_image: np.ndarray) -> Tuple[Optional[str], float]:
        """识别人脸"""
        if not self.known_faces:
            return None, 0.0
        
        # 提取当前人脸特征
        current_features = self.extract_face_features(face_image)
        if current_features is None:
            return None, 0.0
        
        best_match_name = None
        best_similarity = 0.0
        
        # 与所有已知人脸比较
        for employee_id, face_data in self.known_faces.items():
            if face_data['features'] is not None:
                similarity = self.compare_faces(current_features, face_data['features'])
                
                if similarity > best_similarity and similarity > self.recognition_threshold:
                    best_similarity = similarity
                    best_match_name = face_data['name']
        
        return best_match_name, best_similarity
    
    def add_known_face(self, employee_id: str, name: str, face_image: np.ndarray, 
                      features: np.ndarray):
        """添加已知人脸"""
        self.known_faces[employee_id] = {
            'name': name,
            'image': face_image,
            'features': features
        }
    
    def load_known_faces(self, employees_data: List[Tuple]):
        """加载已知人脸数据"""
        self.known_faces = {}
        for employee_data in employees_data:
            # 处理不同长度的数据（兼容4个字段和5个字段）
            if len(employee_data) == 4:
                name, employee_id, face_image, face_features = employee_data
            elif len(employee_data) == 5:
                name, employee_id, face_image, face_features, created_date = employee_data
            else:
                print(f"警告: 员工数据字段数量不正确: {len(employee_data)}")
                continue
            if face_image:
                try:
                    # 解码图像
                    nparr = np.frombuffer(face_image, np.uint8)
                    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
                    
                    # 解码人脸特征
                    features = None
                    if face_features:
                        features = np.frombuffer(face_features, dtype=np.float64)
                    
                    self.known_faces[employee_id] = {
                        'name': name,
                        'image': img,
                        'features': features
                    }
                except Exception as e:
                    print(f"加载人脸数据失败: {e}")
        
        print(f"已加载 {len(self.known_faces)} 个人脸数据")
    
    def get_known_faces(self) -> Dict:
        """获取已知人脸数据"""
        return self.known_faces
    
    def get_employee_id_by_name(self, name: str) -> Optional[str]:
        """根据姓名获取员工ID"""
        for employee_id, face_data in self.known_faces.items():
            if face_data['name'] == name:
                return employee_id
        return None
